import tensorflow as tf
from lstm import BNLSTMCell, LSTMStateTuple
from initializer import get_zero_shape, get_normal_shape

# Model Configuration
class Config(object):
    batch_size = 128
    num_steps = 180
    hidden_size = 2000
    num_layers = 1
    num_epochs = 10
    lr = 0.001
    max_grad_val = 1.
    keep_prob = 1.
def get_config():
    return Config()

# Input info of feeding data
class Text8Input(object):
    def __init__(self, config, data, name=None):
        self.batch_size = batch_size = config.batch_size
        self.num_steps = num_steps = config.num_steps
        self.epoch_size = ((len(data) // batch_size) - 1) // num_steps
        # shape: (batch_size, num_steps)
        # targets is next batch data
        self.input_data, self.targets = _producer(
            data, batch_size, num_steps, name=name)
def _producer(raw_data, batch_size, num_steps, name=None):
    """Iterate on the raw Text8 data.
    This chunks up raw_data into batches of examples and returns Tensors that
    are drawn from these batches.
    Args:
        raw_data: one of the raw data outputs from ptb_raw_data.
        batch_size: int, the batch size.
        num_steps: int, the number of unrolls.
        name: the name of this operation (optional).
    Returns:
        A pair of Tensors, each shaped [batch_size, num_steps]. The second element
        of the tuple is the same data time-shifted to the right by one.
    Raises:
        tf.errors.InvalidArgumentError: if batch_size or num_steps are too high.
    """
    with tf.name_scope(name, "Text8Producer", [raw_data, batch_size, num_steps]):
        raw_data = tf.convert_to_tensor(raw_data, name="raw_data", dtype=tf.int32)

        data_len = tf.size(raw_data)
        batch_len = data_len // batch_size
        data = tf.reshape(raw_data[0 : batch_size * batch_len],
                                            [batch_size, batch_len])
        # The last batch is only used for prediction
        epoch_size = (batch_len - 1) // num_steps
        assertion = tf.assert_positive(
                epoch_size,
                message="epoch_size == 0, decrease batch_size or num_steps")
        with tf.control_dependencies([assertion]):
            epoch_size = tf.identity(epoch_size, name="epoch_size")

        i = tf.train.range_input_producer(epoch_size, shuffle=False).dequeue()
        x = tf.strided_slice(data, [0, i * num_steps],
                                                 [batch_size, (i + 1) * num_steps])
        x.set_shape([batch_size, num_steps])
        y = tf.strided_slice(data, [0, i * num_steps + 1],
                                                 [batch_size, (i + 1) * num_steps + 1])
        y.set_shape([batch_size, num_steps])
        return x, y

# Model Construction
class Text8Model(object):
    def __init__(self, is_training, config, input_, mode):
        self._input = input_
        self._training = is_training

        batch_size = input_.batch_size
        num_steps = input_.num_steps
        size = config.hidden_size
        vocab_size = config.vocab_size

        def lstm_cell():
            return BNLSTMCell(size, num_steps,
                is_training, mode=mode)

        attn_cell = lstm_cell
        if is_training and config.keep_prob < 1:
            def attn_cell():
                return tf.contrib.rnn.DropoutWrapper(
                        lstm_cell(), output_keep_prob=config.keep_prob)
        cell = tf.contrib.rnn.MultiRNNCell(
                [attn_cell() for _ in range(config.num_layers)], state_is_tuple=True)

        self._initial_state = tuple(
            LSTMStateTuple(tf.constant(0, tf.int32),
                get_zero_shape(size, batch_size, tf.float32),
                get_normal_shape(size, batch_size, tf.float32))
            for i in range(config.num_layers))

        with tf.device("/cpu:0"):
            embedding = tf.get_variable(
                    "embedding", [vocab_size, size], dtype=tf.float32)
            inputs = tf.nn.embedding_lookup(embedding, input_.input_data)
        # inputs shape: (B, T, D)

        if is_training and config.keep_prob < 1:
            inputs = tf.nn.dropout(inputs, config.keep_prob)

        outputs, state = tf.nn.dynamic_rnn(cell, inputs,
            initial_state=self._initial_state, dtype=tf.float32)
        output = tf.reshape(outputs, [-1, size])

        softmax_w = tf.get_variable(
                "softmax_w", [size, vocab_size], dtype=tf.float32)
        softmax_b = tf.get_variable("softmax_b", [vocab_size], dtype=tf.float32)
        logits = tf.matmul(output, softmax_w) + softmax_b

        # Reshape logits to be 3-D tensor for sequence loss
        logits = tf.reshape(logits, [batch_size, num_steps, vocab_size])

        # use the contrib sequence loss and average over the batches
        loss = tf.contrib.seq2seq.sequence_loss(
                logits,
                input_.targets,
                tf.ones([batch_size, num_steps], dtype=tf.float32),
                average_across_timesteps=False,
                average_across_batch=True
        )
        # loss shape: (T)

        correct_prediction = tf.equal(
            tf.argmax(tf.nn.softmax(logits), 2),
            tf.cast(input_.targets, tf.int64))
        accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

        # update the cost variables
        # To conclude, loss avg on B, sum on T
        self._cost = cost = tf.reduce_sum(loss)
        self._acc = accuracy
        self._final_state = state

        if not is_training:
            return

        tvars = tf.trainable_variables()
        grads = [None if grad is None else
            tf.clip_by_value(grad, -config.max_grad_val, config.max_grad_val)
            for grad in tf.gradients(cost, tvars)]

        optimizer = tf.train.AdamOptimizer(learning_rate=config.lr)
        self._train_op = optimizer.apply_gradients(
            zip(grads, tvars),
            global_step=tf.contrib.framework.get_or_create_global_step())

    @property
    def input(self):
        return self._input

    @property
    def initial_state(self):
        return self._initial_state

    @property
    def cost(self):
        return self._cost

    @property
    def acc(self):
        return self._acc

    @property
    def final_state(self):
        return self._final_state

    @property
    def training(self):
        return self._training

    @property
    def train_op(self):
        return self._train_op